package com.hawk.eye.client.tracer;

import com.hawk.eye.common.utils.JsonUtil;

import javax.servlet.ReadListener;
import javax.servlet.ServletException;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.Part;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * @className: ContentCachingRequestWrapper
 * @description: 重写HttpServletRequestWrapper，用于读取请求参数
 * @author: salad
 * @date: 2023/9/21
 **/
public class ContentCachingRequestWrapper extends HttpServletRequestWrapper {

    private final byte[] body;

    private final Map<String, Object> params;

    private final Map<String, Object> parts;

    public ContentCachingRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        parts = getMultipartRequest(request);
        if (parts.isEmpty()) {
            body = getRequestBody(super.getInputStream()).getBytes(StandardCharsets.UTF_8);
        } else {
            body = new byte[]{};
        }
        params = getFormParams(request);
    }

    private Map<String, Object> getMultipartRequest(HttpServletRequest request) {
        Map<String, Object> partsMap = new HashMap<>(16);
        try {
            Collection<Part> parts = request.getParts();
            for (Part part : parts) {
                // 处理每个部分，比如获取文件名
                String fileName = part.getSubmittedFileName();
                if (fileName != null) {
                    partsMap.put(fileName, part);
                } else {
                    // 如果是普通字段，记录字段值
                    String formValue = request.getParameter(part.getName());
                    partsMap.put(part.getName(), formValue);
                }
            }
        } catch (IOException | ServletException e) {
            return partsMap;
        }
        return partsMap;
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {

        final ByteArrayInputStream inputStream = new ByteArrayInputStream(body);

        return new ServletInputStream() {

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }

            @Override
            public int read() throws IOException {
                return inputStream.read();
            }
        };
    }

    public Map<String, Object> getParams() {
        return params;
    }

    public Map<String, Object> getMultiParts() {
        return parts;
    }

    public String getMultiPartsAsString() {
        return JsonUtil.toJson(getMultiParts());
    }

    /**
     * 获取GET FORM表单参数
     *
     * @param req
     * @return
     */
    public Map<String, Object> getFormParams(HttpServletRequest req) {
        Map<String, Object> map = new HashMap<>(16);
        try {
            Enumeration<String> paramNames = req.getParameterNames();
            while (paramNames.hasMoreElements()) {
                String paramName = paramNames.nextElement();
                String[] paramValues = req.getParameterValues(paramName);
                if (paramValues.length > 0) {
                    String paramValue = paramValues[0];
                    if (paramValue.length() != 0) {
                        map.put(paramName, paramValue);
                    }
                }
            }
        } catch (Exception e) {
            return map;
        }
        return map;
    }

    public byte[] getBody() {
        return body;
    }


    /**
     * 获取 request 请求体
     *
     * @param servletInputStream servletInputStream
     * @return body
     */
    public static String getRequestBody(ServletInputStream servletInputStream) {
        StringBuilder sb = new StringBuilder();
        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(servletInputStream, StandardCharsets.UTF_8));
            String line;
            while ((line = reader.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (servletInputStream != null) {
                try {
                    servletInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (reader != null) {
                try {
                    reader.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return sb.toString();
    }

}